feat: add support for long-context documents by eurekayuan · Pull Request #179 · NVIDIA-NeMo/Anonymizer

eurekayuan · 2026-06-03T18:42:18Z

Summary

Several stages embedded the whole document in a single prompt and hit DataDesigner's 512K (MAX_RENDERED_LEN) render cap, failing outright on long inputs. Every such stage is now windowed: each chunked generator renders its own per-window prompt and calls the model directly, bypassing the cap. Stages keep a single-call fast path when the rendered prompt already fits, so short-document behavior is unchanged.

Per-stage windowing

Detection (chunked_detection.py, new): Overlapping fixed-size character windows; each window is a raw text slice sent to the detector. Per-window offsets are rebased to global, boundary-touching spans are dropped, and overlaps are resolved (resolve_overlaps).
Validation (chunked_validation.py): Not a text window — batches candidate entities (≤100 per call), each with a ±500-character excerpt. Calls run in parallel across the validator pool with round-robin + failover. Decisions are merged per row; the row is dropped only if every pool member fails.
Augmentation (chunked_augmentation.py): Overlapping character windows over tagged text plus seed JSON. A window dynamically shrinks if its rendered prompt exceeds the cap. Outputs are unioned and deduped by (value, label).
Latent (chunked_latent.py): Same mechanism as augmentation (rewrite mode only); deduped by (label, value).
Substitution map (chunked_replace.py): Abutting newline-aligned windows, no overlap. Each chunk carries the accumulated replacement map and a rolling summary, proposing replacements only for new entities so mappings stay consistent across chunks.
Rewrite generation (chunked_rewrite.py): Abutting newline-aligned windows, no overlap. Runs sequentially, passing a continuity preamble and rolling summary between chunks; rewritten parts are stitched.
Final judge (chunked_final_judge.py, new): Splits original and rewritten text into N positionally-paired slices, scores each, and aggregates per-dimension by minimum. Rubric scales are embedded in the prompt with structured output. Replaces the non-windowedLLMJudgeColumnConfig.

Parallel processing

Stateless stages (detection, validation, augmentation, latent, judge) dispatch windows in parallel (bounded ThreadPoolExecutor; the per-alias rate limit still governs real in-flight calls) and merge afterward.
Stateful stages (substitution-map, rewrite generation) stay sequential to thread the map / rolling summary across seams for consistency.

Window sizing

detection_window_max_render_chars (default 128 KiB, clamped ≤ NDD's render cap) is the single knob; it is threaded into detection, augmentation, latent, substitute-map, rewrite, and judge.
detection_window_safety_margin_chars (8K) leaves headroom for prompt scaffolding; detection_window_overlap_chars (1K) sets the overlap for the overlapping stages; a 4K floor prevents pathological shrinking.

Fault tolerance & failure tracking

Augmentation, latent, and the final judge are resilient to a single bad window: a window whose call fails is logged and skipped rather than dropping the whole record. Skipped-window counts are surfaced in trace_dataframe (COL_AUGMENTATION_FAILED_WINDOWS / COL_LATENT_FAILED_WINDOWS); the judge degrades to defaults if all windows fail.
Detection windows re-raise (the record fails) to preserve detection completeness, and validation relies on pool failover.

Observability

Per-window debug logging across all chunked stages: window ranges/sizes, rendered length vs cap, shrink events, rolling-summary contents, and per-stage entity/replacement/window counts.

Type of Change

Testing

make test passes locally
make check passes locally (format + lint + typecheck + lock-check)
Added/updated tests for changes

Documentation

If docs changed: make docs-build passes locally

github-actions · 2026-06-03T18:42:34Z

All contributors have signed the DCO ✍️ ✅
_{Posted by the DCO Assistant Lite bot.}

eurekayuan · 2026-06-03T18:47:19Z

I have read the DCO document and I hereby sign the DCO.

greptile-apps · 2026-06-03T18:55:41Z

Greptile Summary

This PR adds long-context support across all LLM-heavy stages of the anonymization pipeline by replacing single-call LLMStructuredColumnConfig / LLMTextColumnConfig uses with per-window prompt rendering that bypasses DataDesigner's 512K render cap. Each stage gets a fast path for documents that already fit and a windowed path for longer inputs, with parallel dispatch for stateless stages (detection, augmentation, latent, judge) and sequential context carry-over for stateful ones (substitution map, rewrite generation).

New chunked stages (chunked_detection.py, chunked_augmentation.py, chunked_latent.py, chunked_replace.py, chunked_rewrite.py, chunked_final_judge.py, chunked_steps.py): each implements window planning, LLM dispatch, and a stage-specific merge strategy (offset rebasing + resolve_overlaps for detection; union/dedupe for augmentation/latent/judge; rolling summary + accumulated map for replace/rewrite).
Config knobs (detection_window_max_render_chars, detection_window_safety_margin_chars, detection_window_overlap_chars): single source of truth in Detect config with a model validator that prevents pathological stride collapse; values are threaded through anonymizer.py into all workflows.
Fault tolerance: stateless parallel stages skip failed windows (logging counts in trace_dataframe) rather than failing the record; stateful sequential stages re-raise to preserve correctness guarantees.

Confidence Score: 4/5

Safe to merge with one gap to address: standalone evaluate() doesn't forward custom window sizing to the final judge.

The windowing logic is well-designed throughout — offset rebasing, boundary-span dropping, shrink-on-overflow, and rolling-summary carry-over all work correctly. The one concrete gap is that Anonymizer.evaluate() calls _run_final_judge() without forwarding window_max_render_chars / window_safety_margin_chars, so users who configured a smaller cap to match their LLM's context window won't have it respected during standalone evaluation. Everything in the run() / preview() paths is wired correctly.

Files Needing Attention: src/anonymizer/interface/anonymizer.py (evaluate() call site) and src/anonymizer/engine/replace/chunked_replace.py (new_chunk_entities vs chunk_tagged_text on hard-cut boundaries).

Important Files Changed

Filename	Overview
src/anonymizer/interface/anonymizer.py	Window config threaded through run()/preview() correctly, but evaluate() calls _run_final_judge() without window sizing, silently ignoring custom detection_window_max_render_chars for standalone evaluation.
src/anonymizer/engine/replace/chunked_replace.py	New chunked substitution-map generator with rolling summary; entities straddling hard-cut boundaries can appear in the entity list without tagged context, potentially yielding lower-quality replacements.
src/anonymizer/engine/rewrite/chunked_steps.py	Generic windowed metadata step; windows dispatch serially despite being independent (unlike parallel augmentation/detection stages), with no comment explaining the design choice.
src/anonymizer/engine/rewrite/chunked_rewrite.py	New chunked rewrite with rolling-summary continuity; fast-path measurement includes the continuity preamble that the actual LLM call omits (flagged in prior thread).
src/anonymizer/engine/rewrite/chunked_final_judge.py	New windowed final judge; positional slice pairing, worst-window aggregation, and all-windows-failed fallback are logically correct.
src/anonymizer/engine/detection/chunked_detection.py	New overlapping-window seed detector; offset rebasing, edge-span dropping, and resolve_overlaps merge are logically sound.
src/anonymizer/engine/detection/chunked_augmentation.py	New windowed augmentation with shrink-on-overflow loop and parallel dispatch; fast path and per-window failure resilience are correct.
src/anonymizer/engine/windowing.py	New boundary-aware windowing utility; delimiter fallback logic and MIN_WINDOW_FRACTION guard are correct and well-tested.
src/anonymizer/config/anonymizer_config.py	Adds three window-sizing fields to Detect config; the model_validator that rejects overlap >= effective_window is a good guard against pathological stride collapse.
src/anonymizer/engine/rewrite/qa_generation.py	New batched quality-QA path; fast path and batched path both use model_dump() without mode="json" (flagged in prior thread); imports private _compile_template from chunked_steps (also prior thread).
src/anonymizer/engine/rewrite/final_judge.py	Replaces LLMJudgeColumnConfig with windowed judge; rubric scales correctly embedded in prompt, template variables updated to original_text/rewritten_text to match render_judge_prompt.
src/anonymizer/engine/rewrite/sensitivity_disposition.py	Converts to windowed disposition; merge correctly picks most-protective entry per entity across windows; full entity list passed to every window is intentional and documented.
src/anonymizer/engine/rewrite/domain_classification.py	Converts to windowed classification with first_only=True; _first_output has an unguarded outputs[0] access (flagged in prior thread).

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Input Document] --> B{len <= cap?}
    B -- Yes / Fast path --> C[Single LLM call]
    B -- No / Windowed path --> D[Plan Windows]
    D --> E[Detection\nchunked_detection.py\noverlapping + parallel]
    D --> F[Augmentation\nchunked_augmentation.py\noverlapping + shrink + parallel]
    D --> G[Latent\nchunked_latent.py\noverlapping + parallel]
    D --> H[Disposition / Domain / QA\nchunked_steps.py\nboundary-aligned + serial]
    D --> I[Substitute Map\nchunked_replace.py\nboundary-aligned + rolling summary]
    D --> J[Rewrite\nchunked_rewrite.py\nboundary-aligned + rolling summary]
    D --> K[Final Judge\nchunked_final_judge.py\npositional slices + parallel]
    E --> L[rebase + drop edges + resolve_overlaps]
    F --> M[union dedupe by value+label]
    G --> M
    H --> N[stage-specific merge fn]
    I --> O[accumulated map merge_replacements]
    J --> P[stitch parts]
    K --> Q[min score per dimension]
    L & M & N & O & P & Q --> R[Write output column]
    C --> R

Comments Outside Diff (1)

src/anonymizer/interface/anonymizer.py, line 400-405 (link)

Custom window config ignored in standalone evaluate()

_run_final_judge() is called without window_max_render_chars / window_safety_margin_chars, so the call always falls back to the 128 KiB module-level defaults — regardless of what the user configured via detection_window_max_render_chars. For users who lowered the cap specifically to stay within their LLM's context window or avoid timeouts, the standalone evaluate path silently uses a larger window than intended. Concrete failure scenario: a user sets detection_window_max_render_chars=65536 to match a 64 K-token model, their documents are ~100 KB, run() respects the 64 KB cap and windows correctly, but evaluate() uses the 128 KB default and submits a single over-budget prompt, causing the judge call to time out. Since AnonymizerResult doesn't carry the detect config, the window sizing would need to be threaded through EvaluateConfig (currently a placeholder).

_{Reviews (3): Last reviewed commit: "Windowing: require boundaries in the win..." | Re-trigger Greptile}

greptile-apps · 2026-06-03T18:55:45Z

+                _clip(summary),
+            )
+
+    stitched = "\n".join(part for part in rewritten_parts if part)


Chunk boundaries are aligned to newlines by iter_boundary_windows, so each tagged[start:end] slice already ends with " ". When the LLM mirrors that structure in its output (natural for paragraph-aware models), every rewritten_chunk also ends with " ", and " ".join(...) then inserts a second newline — producing a blank line between every chunk boundary in the final anonymized document. Joining with "" is sufficient because the delimiter is already part of each chunk.

Suggested change

stitched = "\n".join(part for part in rewritten_parts if part)

stitched = "".join(part for part in rewritten_parts if part)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

greptile-apps · 2026-06-03T18:55:46Z

+    # Fast path: the full single-call rewrite prompt fits under the cap.
+    single_rendered = _render_chunk_prompt(template=params.single_call_prompt_template, chunk_row=row, summary="")
+    if len(single_rendered) <= cap:
+        logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap)
+        text = _rewrite_chunk(
+            facade=facade,
+            prompt=_compile_template(params.single_call_prompt_template).render(**row),
+            system_prompt=params.system_prompt,
+            purpose="rewrite-generation",
+        )


The fast path measures single_rendered as _render_chunk_prompt(..., summary=""), which prepends the ~270-char continuity preamble, but then the actual LLM call omits that preamble. This means a document whose body-only prompt falls in (cap - 270, cap] chars will be routed into the chunked path unnecessarily. Measure with just the body to match what is actually sent.

Suggested change

# Fast path: the full single-call rewrite prompt fits under the cap.

single_rendered = _render_chunk_prompt(template=params.single_call_prompt_template, chunk_row=row, summary="")

if len(single_rendered) <= cap:

logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap)

text = _rewrite_chunk(

facade=facade,

prompt=_compile_template(params.single_call_prompt_template).render(**row),

system_prompt=params.system_prompt,

purpose="rewrite-generation",

)

# Fast path: measure body-only prompt (no continuity preamble) since that is what is sent.

single_rendered = _compile_template(params.single_call_prompt_template).render(**row)

if len(single_rendered) <= cap:

logger.debug("rewrite: single-call fast path (rendered=%d chars <= cap=%d)", len(single_rendered), cap)

text = _rewrite_chunk(

facade=facade,

prompt=single_rendered,

system_prompt=params.system_prompt,

purpose="rewrite-generation",

)

greptile-apps · 2026-06-03T18:55:47Z

 )
 from anonymizer.engine.ndd.model_loader import resolve_model_alias
 from anonymizer.engine.prompt_utils import substitute_placeholders


Private symbol imported across module boundary. _compile_template is module-private (underscore-prefixed) in chunked_steps.py. Importing it here creates a hidden coupling: if the function is renamed or inlined, qa_generation.py breaks without any clear contract. Consider exposing it as a public helper in chunked_steps.py or defining a local copy with its own lru_cache in this module.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

greptile-apps · 2026-06-03T18:55:48Z

+        row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump()
+        return row
+
+    units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]")
+    base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"}))
+    batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars)
+    items: list[dict[str, Any]] = []
+    for batch_idx, batch in enumerate(batches):
+        rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)})
+        out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}")
+        for item in out.items:
+            items.append({**item.model_dump(mode="json"), "id": len(items) + 1})
+    row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump()


The fast path stores the result via .model_dump() (no mode="json"), while every other windowed generator in this PR consistently uses .model_dump(mode="json"). Without mode="json", Pydantic returns native Python objects rather than JSON-serializable equivalents, which can cause downstream serialization failures. The batched path has the same inconsistency.

Suggested change

row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump()

return row

units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]")

base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"}))

batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars)

items: list[dict[str, Any]] = []

for batch_idx, batch in enumerate(batches):

rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)})

out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}")

for item in out.items:

items.append({**item.model_dump(mode="json"), "id": len(items) + 1})

row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump()

row[COL_QUALITY_QA] = _generate(full_rendered, "quality-qa-generation").model_dump(mode="json")

return row

units = json.loads(row.get(COL_MEANING_UNITS_SERIALIZED) or "[]")

base_len = len(compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: "[]"}))

batches = _batch_units_by_size(units, base_len, max_render_chars - safety_margin_chars)

items: list[dict[str, Any]] = []

for batch_idx, batch in enumerate(batches):

rendered = compiled.render(**{**row, COL_MEANING_UNITS_SERIALIZED: json.dumps(batch, ensure_ascii=False)})

out = _generate(rendered, f"quality-qa-generation-batch-{batch_idx}")

for item in out.items:

items.append({**item.model_dump(mode="json"), "id": len(items) + 1})

row[COL_QUALITY_QA] = QualityQAPairsSchema.model_validate({"items": items}).model_dump(mode="json")

greptile-apps · 2026-06-03T18:55:49Z

+_DEFAULT_MAX_RENDER_CHARS: int = _DetectConfig.model_fields["detection_window_max_render_chars"].default
+_DEFAULT_SAFETY_MARGIN_CHARS: int = _DetectConfig.model_fields["detection_window_safety_margin_chars"].default
+
+


Unguarded index on potentially empty list. _first_output calls outputs[0] without checking length. In run_windowed_step with first_only=True, if iter_boundary_windows returns an empty list, outputs is empty and this raises IndexError. The fast path makes this unreachable today, but a defensive guard would make the failure mode explicit.

andreatnvidia

Thanks for taking this on. This is a substantial first PR, and the overall direction makes sense: split long records into bounded windows, carry forward the state needed for consistency, and keep the replacement map explicit.

I left a few comments on edge cases I think are worth tightening before merge. The main themes are:

thread the user-supplied window sizing through every windowed stage
make per-window failures local where possible, instead of dropping the whole record
validate overlap settings early so a bad config cannot explode into thousands of model calls
avoid silently accepting empty rewrite chunks as successful output

The tests and docs coverage are in good shape, and I think the feature is close. These changes should make it more reliable on real long documents.

andreatnvidia · 2026-06-09T14:13:24Z

            ),
            *self._qa_wf.columns(selected_models=selected_models),
            *self._rewrite_gen_wf.columns(
+                window_max_render_chars=window_max_render_chars,


this only threads the user-supplied window cap into rewrite generation. domain classification, sensitivity disposition, QA generation, and final judge still build their window params from module defaults, so a user who lowers Detect.detection_window_max_render_chars still gets ~128k prompts in those stages. Could pass the same kwargs through those columns() calls and _run_final_judge too?

andreatnvidia · 2026-06-09T14:13:24Z

+    if params.first_only:
+        windows = windows[:1]
+    outputs = []
+    for start, end in windows:


Claude Code caught this one: once this takes the windowed path, a single transient model error or a chunk that legitimately has no meaning units can drop the whole record. Could wrap each window call, skip/log failed windows, and handle the all-failed case explicitly?

andreatnvidia · 2026-06-09T14:13:24Z

+            "prompt scaffolding and tags when sizing augmentation/latent windows."
+        ),
+    )
+    detection_window_overlap_chars: int = Field(


suggestion: can we validate that detection_window_overlap_chars is smaller than the effective window size? Right now overlap == window is accepted and the planners advance one character at a time. My smoke test turned a 20k-char row into 16,001 windows.

andreatnvidia · 2026-06-09T14:13:24Z

+                _clip(summary),
+            )
+
+    stitched = "\n".join(part for part in rewritten_parts if part)


separate from the newline-stitching comment already here: filtering with if part also hides an empty rewrite chunk. If one window returns {"rewritten_text": ""}, that section disappears with no failed-window count or review signal. Maybe count/flag empty chunks instead of treating them as successful output?

…ate overlap, flag empty rewrite chunks Signed-off-by: eurekayuan <zhuoweny@nvidia.com>

…workflow-column plugins Upstream (#182, #190, 8e6eed5) made the detection workflow exportable for external/distributed executors, requiring every column config to be JSON- serializable. Convert the windowed CustomColumnConfig closures introduced by this branch into plugin column configs (WindowedDetectionConfig, WindowedAugmentationConfig, WindowedLatentConfig) registered via data_designer.plugins entry points, mirroring the chunked-validation plugin. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…repeated content The continuity preamble's 'do NOT repeat earlier content' made the rewriter emit near-empty output for sections resembling earlier ones, collapsing repetitive long documents (~300K chars -> ~3-18K). Instruct each window to be rewritten in full so output length tracks input structure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…> period -> warned hard cut A delimiter occurring only near the start of a window previously produced a degenerate, tiny window (many extra LLM calls on documents with sparse newlines). Boundaries now count only if they keep the window at least half of max_chars; when the primary delimiter has no such occurrence the sentence boundary '.' is tried, and if that also fails the window is hard-cut at max_chars with a warning. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

eurekayuan requested review from a team as code owners June 3, 2026 18:42

eurekayuan changed the title ~~Add support for handling long-context docs~~ feat/long-context Jun 3, 2026

eurekayuan changed the title ~~feat/long-context~~ feat: add support for long-context documents Jun 3, 2026

greptile-apps Bot reviewed Jun 3, 2026

View reviewed changes

andreatnvidia reviewed Jun 9, 2026

View reviewed changes

eurekayuan and others added 7 commits July 23, 2026 14:46

Add long-context windowing for detection, replace, and rewrite

7f96979

chunking final judge

c0ee8ea

format scripts

1c0878d

Address review: thread window sizing, localize window failures, valid…

0662e36

…ate overlap, flag empty rewrite chunks Signed-off-by: eurekayuan <zhuoweny@nvidia.com>

eurekayuan force-pushed the feature/long-context branch from 6639ecb to 051f416 Compare July 24, 2026 22:57

	stitched = "\n".join(part for part in rewritten_parts if part)
	stitched = "".join(part for part in rewritten_parts if part)

		_DEFAULT_MAX_RENDER_CHARS: int = _DetectConfig.model_fields["detection_window_max_render_chars"].default
		_DEFAULT_SAFETY_MARGIN_CHARS: int = _DetectConfig.model_fields["detection_window_safety_margin_chars"].default

Uh oh!

Conversation

eurekayuan commented Jun 3, 2026

Summary

Per-stage windowing

Parallel processing

Window sizing

Fault tolerance & failure tracking

Observability

Type of Change

Testing

Documentation

Uh oh!

github-actions Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eurekayuan commented Jun 3, 2026

Uh oh!

greptile-apps Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Flowchart

Comments Outside Diff (1)

Uh oh!

greptile-apps Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

andreatnvidia left a comment

Choose a reason for hiding this comment

Uh oh!

andreatnvidia Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

andreatnvidia Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

andreatnvidia Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

andreatnvidia Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

github-actions Bot commented Jun 3, 2026 •

edited

Loading

greptile-apps Bot commented Jun 3, 2026 •

edited

Loading